import os
from operator import itemgetter

import bs4
from langchain.agents import AgentType
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.history_aware_retriever import create_history_aware_retriever
from langchain.chains.retrieval import create_retrieval_chain
from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_chroma import Chroma
from langchain_community.agent_toolkits import SQLDatabaseToolkit, create_sql_agent

from langchain_community.document_loaders import WebBaseLoader
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_community.tools import TavilySearchResults, QuerySQLDataBaseTool
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, PromptTemplate
from langchain_core.runnables import RunnableWithMessageHistory, RunnablePassthrough
from langchain_openai import AzureChatOpenAI, AzureOpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langgraph.prebuilt import chat_agent_executor

os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "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"

os.environ["AZURE_OPENAI_ENDPOINT"] = "http://menshen.test.xdf.cn"
# os.environ["OPENAI_API_BASE"] = "http://menshen.test.xdf.cn"
os.environ["OPENAI_API_KEY"] = "c8575027653b42b1b47747f0b4ab135b"
os.environ["OPENAI_API_TYPE"] = "azure"
os.environ["OPENAI_API_VERSION"] = "2023-05-15"

llm = AzureChatOpenAI(
    deployment_name="gpt-4o",
    model_name="gpt-4o",
    temperature=0
)

USER_NAME = 'ourcrm_admin'
PWD = 'ourcrm_admin'
HOST = '172.24.30.115'
PORT = 3306
DATABASE = 'scrm_test'

# 使用mysqlclient驱动
mysql_uri = f'mysql+mysqldb://{USER_NAME}:{PWD}@{HOST}:{PORT}/{DATABASE}?charset=utf8mb4'

db = SQLDatabase.from_uri(mysql_uri,include_tables=["crm_sys_user"])
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
print(toolkit.get_tools())

agent_executor = create_sql_agent(llm=llm, toolkit=toolkit, verbose=False, agent_type=AgentType.OPENAI_FUNCTIONS)

resp = agent_executor.invoke("crm_sys_user表中有多少条数据")
print(resp)